import logging
from collections import namedtuple
from pathlib import Path

import coloredlogs
import gradio as gr
import requests
import torch
import yaml
from transformers import AutoModelForCausalLM, AutoTokenizer

from args import parse_args
from models import ModelManager
from utils import (
    check_balance,
    dummy_generate_text,
    format_message_history,
    reset_dialog,
)

log = logging.getLogger(__name__)
args = parse_args(known_args_only=True)
coloredlogs.install(level="DEBUG" if args.debug else "INFO")
config = yaml.safe_load(args.config.read_text())
config = namedtuple("Config", config.keys())(*config.values())
assert config.default_model in config.model_paths, "Default model not in model paths"

COPY_JS = Path("copy.js").read_text()
COPY_HTML = config.copy_html.format(COPY_JS)
with open(config.sys_prompt_file) as f:
    sys_prompts = yaml.safe_load(f)
assert "default" in sys_prompts, "Default system prompt not found"


def generate_text(
    sys_prompt,
    custom_sys_prompt,
    prompt,
    max_new_tokens,
    model,
    enable_hist,
    messages,
    request: gr.Request,
):
    log.debug("All params:")
    log.debug("\n".join(f"\t{k}: {v}" for k, v in locals().items()))
    log.info(
        f"Query by {request.username} < {repr(prompt[:config.terminal_preview_len])}"
    )
    sys_prompt = custom_sys_prompt or sys_prompt
    model_manager.update_model(model)
    if max_new_tokens == 0:
        max_new_tokens = None
    resp, new_msgs = model_manager(
        sys_prompt, prompt, max_new_tokens, messages if enable_hist else []
    )
    if not enable_hist:  # model received no history
        new_msgs = messages + new_msgs
    history = format_message_history(new_msgs)

    # calculate cost
    global prev_balance
    bal = check_balance(config)
    cost = None if prev_balance is None or bal is None else prev_balance - bal
    prev_balance = bal
    cost_str = f"Cost: {round(cost, 5) if cost else 'N/A'} | Balance: {bal}"

    log.info(
        f"{model_manager.current.split('/')[-1]} > "
        f"{repr(resp[:config.terminal_preview_len])}"
    )
    log.info(cost_str)
    return resp, history, cost_str, new_msgs


def setup_interface():
    with gr.Blocks(title=config.app_title, js="./copy.js", css="./style.css") as demo:
        # add title
        msg_state = gr.State([])
        gr.Markdown(f"# {config.app_title} v{config.display_version}")
        gr.Markdown(config.app_description)

        with gr.Accordion("Config"):
            with gr.Row():
                # model choice
                model_path = gr.Radio(
                    list(config.model_paths.items()),
                    value=lambda: model_manager.current,
                    label="Model",
                    elem_id="model-choice",
                )
                # sys prompt choice
                sys_prompt = gr.Radio(
                    list(sys_prompts.items()),
                    value=sys_prompts["default"],
                    label="System Prompt",
                )
            custom_sys_prompt = gr.Textbox(
                label="Custom System Prompt",
                lines=1,
                placeholder="Leave empty to use predefined prompts",
            )
            max_new_token = gr.Slider(
                0,
                config.max_new_token_maxval,
                0,
                label="Max new tokens (0 for unlimited; not supported for all models)",
            )

        # main interface
        history = gr.Chatbot(
            label="Session History",
            elem_id="history",
            placeholder="No session history yet",
        )
        prompt = gr.Textbox(
            label="Your Prayer",
            lines=config.prompt_lines,
            placeholder="Enter your prompt here",
        )

        # buttons
        with gr.Row():
            enable_hist = gr.Checkbox(
                False,
                label="Remember my prayers (progressively costly)",
                elem_id="hist-cb",
            )
            reset = gr.Button(
                "Forget", variant="stop", elem_id="reset-btn", interactive=False
            )
            submit = gr.Button("Invoke", variant="primary", elem_id="submit-btn")
            gr.Button("Transcribe", variant="secondary", elem_id="copy-btn")

        # logs
        with gr.Accordion("Info"):
            balance = gr.Textbox(
                label="The Toll",
                placeholder="Start a session to check balance",
                elem_id="balance",
            )
            output = gr.Textbox(
                label="Their Blessing", lines=config.prompt_lines, elem_id="gen"
            )

        def toggle_history(enabled):
            new_reset = gr.Button(
                "Forget", variant="stop", elem_id="reset-btn", interactive=enabled
            )
            return new_reset

        enable_hist.input(
            toggle_history,
            inputs=enable_hist,
            outputs=reset,
        )
        reset.click(reset_dialog, inputs=msg_state, outputs=[msg_state, history])
        submit.click(
            fn=generate_text if not args.debug else dummy_generate_text,
            inputs=[
                sys_prompt,
                custom_sys_prompt,
                prompt,
                max_new_token,
                model_path,
                enable_hist,
                msg_state,
            ],
            outputs=[output, history, balance, msg_state],
        )

    return demo


def simple_auth(user, passwd):
    auth_file = yaml.safe_load(Path(config.auth_file).read_text())
    return user in auth_file and passwd == auth_file[user]


prev_balance = check_balance(config)
model_manager = ModelManager(args, config)
demo = setup_interface()
demo.launch(
    server_name=config.host,
    server_port=config.port,
    auth=None if args.no_auth else simple_auth,
    favicon_path=config.favicon_file,
)
